# encoding: utf-8

''' 任务训练器 '''


import numpy as np
from .aug import mean_std
from model.multipoly_regression import train_multi
from visual.visual import imshow


# 第一题
# 总表训练器
def m_trainer(X: np.ndarray, Y: np.ndarray, iters):
    for year in iters:
        t, tr = train_multi(X, Y)
        pr = t.predict(tr.fit_transform(np.array([year])))
        m,s = mean_std(Y)
        X = np.append(X, [year], axis=0)
        Y = np.append(Y, pr, axis=0)
    # print(t.predict(tr.fit_transform(np.array([[2019]]))))
    # imshow(X, Y, X,t.predict(tr.fit_transform(X)), xlabel='years', ylabel='Account')
    return t.predict(tr.fit_transform(X))
    # imshow(X, Y, X, t.predict(X), xlabel='years', ylabel='Account')

def all_trainer(trainer, X: np.ndarray, Y: np.ndarray, iters=[2017, 2018, 2019]):
    for year in iters:
        t = trainer(X, Y)
        pr = t.predict(np.array([year]))
        m,s = mean_std(Y)
        # print(m,s);exit()
        pr += np.random.normal(m/100, s/100)
        X = np.append(X, [year], axis=0)
        Y = np.append(Y, pr, axis=0)
    res = t.predict(X)
    # imshow(X, Y, X, res, xlabel='years', ylabel='Account')
    return res
